/*
编写者：王出日
日期：2025，1，20
版本 0.4.2
功能：Agent类，实现了最简单的工具执行智能体的功能。
包含：
Agent类：
*/

import 'dart:async';
import 'dart:convert';
import 'BaseAPI.dart';
import 'Tools.dart';

/// Agent响应结果类
class AgentResponse {
  final String role;
  final String content;
  final String? reasoningContent;
  final String? toolName;
  final String? toolArguments;
  
  AgentResponse({
    required this.role,
    required this.content,
    this.reasoningContent,
    this.toolName,
    this.toolArguments,
  });
  
  Map<String, dynamic> toMap() {
    final result = <String, dynamic>{
      'role': role,
      'content': content,
    };
    
    if (reasoningContent != null) {
      result['reasoning_content'] = reasoningContent;
    }
    if (toolName != null) {
      result['tool_name'] = toolName;
    }
    if (toolArguments != null) {
      result['tool_arguments'] = toolArguments;
    }
    
    return result;
  }
}

/// 工具调用结果记录
class ToolCallResult {
  final String id;
  final String name;
  final String arguments;
  final String result;
  final bool success;
  
  ToolCallResult({
    required this.id,
    required this.name,
    required this.arguments,
    required this.result,
    required this.success,
  });
}

/// 最简单的工具执行智能体
class Agent {
  final BaseAPI llm;
  final Tools tools;
  final List<ToolCallResult> toolsCallResult = [];
  final List<Map<String, dynamic>> toolsCall = [];
  final bool isExecute;
  final int contextLimit;
  
  late final String prompt;
  final List<Map<String, dynamic>> messages = [];

  Agent({
    required this.llm,
    required this.tools,
    String? sysPrompt,
    this.isExecute = true,
    this.contextLimit = 8000,
  }) {
    if (sysPrompt != null) {
      prompt = sysPrompt;
      messages.addAll([
        {"role": "system", "content": prompt},
        {"role": "system", "content": ""}, // 目标设置
        {"role": "system", "content": ""}, // 目标达成情况和下一步
      ]);
    } else {
      prompt = "你是一个有用的AI助手，可以使用工具来帮助用户解决问题。";
      messages.addAll([
        {"role": "system", "content": prompt},
        {"role": "system", "content": ""}, // 目标设置
        {"role": "system", "content": ""}, // 目标达成情况和下一步
      ]);
    }
  }

  /// 禁用工具
  bool disableTool(String toolName) {
    return tools.disableTool(toolName);
  }

  /// 启用工具
  bool enableTool(String toolName) {
    return tools.enableTool(toolName);
  }

  /// 获取当前Agent的消息列表
  List<Map<String, dynamic>> getMessages() {
    return List.from(messages);
  }

  /// 清理当前Agent的消息列表，只保留前三个系统消息
  void clearMessages() {
    final systemMessages = messages.take(3).toList();
    messages.clear();
    messages.addAll(systemMessages);
  }

  /// 获取当前Agent的工具列表
  List<Map<String, dynamic>> getTools() {
    return tools.getTools();
  }

  /// 获取当前Agent的提示词
  String getPrompt() {
    return prompt;
  }

  /// 设置当前Agent的目标（更新第二个系统消息）
  void setGoal(String goal) {
    if (messages.length > 1) {
      messages[1] = {"role": "system", "content": goal};
    }
  }

  /// 更新目标达成情况和下一步（更新第三个系统消息）
  void updateGoalStatus(String status) {
    if (messages.length > 2) {
      messages[2] = {"role": "system", "content": status};
    }
  }

  /// 取消当前Agent的目标（清空第二个系统消息）
  void cancelGoal() {
    if (messages.length > 1) {
      messages[1] = {"role": "system", "content": ""};
    }
  }

  /// 在当前的Agent添加新的消息
  void addMessage({
    String? role,
    String? content,
    List<Map<String, dynamic>>? messageList,
  }) {
    if (contextLimit > 0) {
      _limitMessages();
    }

    if (messageList != null) {
      messages.addAll(messageList);
      return;
    }

    if (role != null && content != null) {
      messages.add({"role": role, "content": content});
    }
  }

  /// 获取当前Agent的工具调用结果列表
  List<ToolCallResult> getToolsCallResult() {
    return List.from(toolsCallResult);
  }

  /// 获取当前Agent的工具调用列表
  List<Map<String, dynamic>> getToolsCall() {
    return List.from(toolsCall);
  }

  /// 限制消息列表总字数不超过contextLimit
  void _limitMessages() {
    if (contextLimit < 0) return;
    if (messages.isEmpty) return;

    // 强制保留前三个系统消息
    final preserved = messages.take(3).toList();
    int currentLength = preserved
        .map((msg) => (msg['content'] as String? ?? '').length)
        .fold(0, (sum, len) => sum + len);

    // 如果初始长度已超限，只保留前三个
    if (currentLength >= contextLimit) {
      messages.clear();
      messages.addAll(preserved);
      return;
    }

    // 保留足够多的最新消息
    final remaining = contextLimit - currentLength;
    final newMessages = <Map<String, dynamic>>[];
    newMessages.addAll(preserved);
    
    int total = 0;
    final otherMessages = messages.skip(3).toList().reversed.toList();
    
    for (final msg in otherMessages) {
      final contentLen = (msg['content'] as String? ?? '').length;
      if (total + contentLen <= remaining) {
        newMessages.add(msg);
        total += contentLen;
      } else {
        break;
      }
    }

    // 恢复消息顺序（除了前三个系统消息）
    final finalMessages = <Map<String, dynamic>>[];
    finalMessages.addAll(preserved);
    finalMessages.addAll(newMessages.skip(3).toList().reversed);
    
    messages.clear();
    messages.addAll(finalMessages);
  }

  /// 调用agent进行生成文本回复，非流式输出
  Future<Map<String, dynamic>> predict({
    String? inputText,
    double temperature = 0.5,
    double topP = 0.9,
  }) async {
    if (inputText != null) {
      messages.add({"role": "user", "content": inputText});
    }

    while (true) {
      final llmResult = await llm.predict(
        messages: messages,
        temperature: temperature,
        tools: tools.getTools(),
        topP: topP,
        stream: false,
      ) as Map<String, dynamic>;

      if (llmResult.containsKey('tool_calls')) {
        final toolCalls = llmResult['tool_calls'] as Map<String, dynamic>;
        final functionInfo = toolCalls['function'] as Map<String, dynamic>;
        final toolName = functionInfo['name'] as String;
        final arguments = json.decode(functionInfo['arguments'] as String);

        final result = await _executeAndEnterToGoalLoop(
          toolName,
          arguments,
          toolCalls,
        );

        if (!result['continue'] as bool) {
          return result['response'] as Map<String, dynamic>;
        }

        messages.add({
          "role": "system",
          "content": "工具的执行结果为：\n${result['output']}"
        });
        continue;
      } else {
        messages.add(llmResult);
        return llmResult;
      }
    }
  }

  /// 调用agent进行生成文本回复，流式输出
  Stream<AgentResponse> predictStream({
    String? inputText,
    double temperature = 0.5,
    double topP = 0.9,
  }) async* {
    if (inputText != null) {
      messages.add({"role": "user", "content": inputText});
    }

    final llmResult = await llm.predict(
      messages: messages,
      temperature: temperature,
      tools: tools.getTools(),
      topP: topP,
      stream: true,
    );
    
    if (llmResult is Stream<Map<String, dynamic>>) {
      yield* _parseStream(llmResult);
    } else {
      throw Exception('Expected stream response but got non-stream');
    }
  }

  /// 解析流式响应
  Stream<AgentResponse> _parseStream(
    Stream<Map<String, dynamic>> llmStream,
  ) async* {
    final contentParts = <String>[];
    String reasoningBuffer = "";

    await for (final chunk in llmStream) {
      // 处理工具名称
      if (chunk.containsKey('tool_name')) {
        yield AgentResponse(
          role: 'assistant',
          content: '',
          toolName: chunk['tool_name'] as String,
        );
      }
      // 处理工具调用
      else if (chunk.containsKey('tool_calls') && chunk['id'] != '') {
        final wholeContent = contentParts.join('');
        if (wholeContent.isNotEmpty) {
          messages.add({"role": "assistant", "content": wholeContent});
        }
        contentParts.clear();
        reasoningBuffer = "";

        final toolCalls = (chunk['tool_calls'] as List).first as Map<String, dynamic>;
        final functionInfo = toolCalls['function'] as Map<String, dynamic>;
        final functionArgs = functionInfo['arguments'] as String;

        yield AgentResponse(
          role: 'assistant',
          content: '',
          toolArguments: functionArgs,
        );

        try {
          final args = json.decode(functionArgs) as Map<String, dynamic>;
          
          final toolCallObj = {
            "id": chunk["id"],
            "function": {
              "name": functionInfo["name"],
              "arguments": functionArgs
            }
          };

          toolsCall.add(toolCallObj);

          if (isExecute) {
            final toolResult = await _executeAndEnterToGoalLoop(
              functionInfo["name"] as String,
              args,
              toolCallObj,
            );

            yield AgentResponse(
              role: 'tool',
              content: toolResult['output'] as String,
            );

            if (toolResult['continue'] as bool) {
              messages.add({"role": "assistant", "tool_calls": [toolCallObj]});
              messages.add({
                "role": "tool",
                "content": "工具调用结果：\n${toolResult['output']}"
              });

              // 递归调用生成新的响应
              await for (final newResponse in _parseStream(
                await llm.predict(
                  messages: messages,
                  temperature: 0.5,
                  tools: tools.getTools(),
                  topP: 0.9,
                  stream: true,
                ) as Stream<Map<String, dynamic>>
              )) {
                yield newResponse;
              }
            }
          }
        } catch (e) {
          final errorMsg = "工具参数解析失败: $e";
          yield AgentResponse(role: 'tool', content: errorMsg);
          
          await for (final newResponse in _parseStream(
            await llm.predict(
              inputText: "$errorMsg，请重新输入",
              stream: true,
            ) as Stream<Map<String, dynamic>>
          )) {
            yield newResponse;
          }
        }
      }
      // 处理推理内容
      else if (chunk.containsKey('reasoning_content')) {
        final reasoningContent = chunk['reasoning_content'] as String? ?? '';
        if (reasoningContent.isNotEmpty) {
          reasoningBuffer += reasoningContent;
          yield AgentResponse(
            role: 'assistant',
            content: '',
            reasoningContent: reasoningContent,
          );
        }
      }
      // 处理普通内容
      else {
        final content = chunk['content'] as String? ?? '';
        if (content.isNotEmpty) {
          contentParts.add(content);
          yield AgentResponse(role: 'assistant', content: content);
        }
      }
    }

    // 处理最后的内容
    final wholeContent = contentParts.join('');
    if (wholeContent.isNotEmpty) {
      messages.add({"role": "assistant", "content": wholeContent});
    }
  }

  /// 执行工具调用并进入目标循环
  Future<Map<String, dynamic>> _executeAndEnterToGoalLoop(
    String toolName,
    Map<String, dynamic> args,
    Map<String, dynamic> toolCallObj,
  ) async {
    String toolResult;
    bool shouldContinue = true;

    // 根据工具名称选择执行方式
    if (toolName == "setGoal") {
      setGoal(args["goal"] as String);
      toolResult = "目标设置成功";
    } else if (toolName == "cancelGoal") {
      cancelGoal();
      toolResult = "目标已取消";
    } else if (toolName == "updateGoalStatus") {
      updateGoalStatus(args["status"] as String);
      toolResult = "目标状态已更新";
    } else {
      // 执行普通工具
      try {
        toolResult = await tools.execute(toolName, kwargs: args);
      } catch (e) {
        toolResult = "工具执行失败: $e";
        shouldContinue = false;
      }
    }

    // 记录工具调用结果
    final resultRecord = ToolCallResult(
      id: toolCallObj["id"] as String,
      name: toolName,
      arguments: toolCallObj["function"]["arguments"] as String,
      result: toolResult,
      success: shouldContinue,
    );

    toolsCallResult.add(resultRecord);

    return {
      'output': toolResult,
      'continue': shouldContinue,
      'response': {"role": "assistant", "content": toolResult},
    };
  }
}
